# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Generate data for test."""
import numpy as np

def get_init_params():
    # Generate initialization parameters
    np.random.seed(42)
    return {
        "t": 0.01 * np.random.randn(1, 1, 2, 50),
        "freqs": 0.01 * np.random.randn(1, 1, 2, 48),
    }

def get_golden() -> dict[str, np.ndarray]:
    """Generate golden data for test."""
    output_1 = np.array(
        [[[[5.06591797e-03, 6.43920898e-03, -2.27355957e-03, 1.57470703e-02,
            -4.69970703e-03, -4.66918945e-03, 2.74658203e-03, -1.72119141e-02,
            -1.01318359e-02, -9.09423828e-03, 1.46484375e-02, 6.71386719e-04,
            -5.43212891e-03, -1.16577148e-02, -6.01196289e-03, -6.04248047e-03,
            -1.32560730e-04, 8.23974609e-03, 2.24304199e-03, -1.33056641e-02,
            7.38525391e-03, -1.22833252e-03, -1.47094727e-02, -4.39453125e-03,
            -1.47247314e-03, 1.52587891e-02, -2.25830078e-03, 7.65991211e-03,
            5.43212891e-03, -4.69970703e-03, -1.92871094e-02, -5.61523438e-03,
            2.99072266e-03, -1.40991211e-02, -2.39562988e-03, -1.42822266e-02,
            1.13677979e-03, 3.75366211e-03, -3.06701660e-03, 1.85546875e-02,
            -1.05590820e-02, -1.23291016e-02, -1.96533203e-02, 1.93023682e-03,
            1.72424316e-03, -2.99072266e-03, -7.08007812e-03, 1.06811523e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.26538086e-03, -6.77490234e-03, 1.03149414e-02, -8.30078125e-03,
            3.28063965e-03, -4.79125977e-03, -1.10473633e-02, 8.17871094e-03,
            -7.89642334e-04, 3.61633301e-03, 3.61633301e-03, -3.35693359e-04,
            -2.62451172e-02, 6.63757324e-04, 9.46044922e-04, -2.19726562e-03,
            1.48315430e-02, -7.99560547e-03, 9.21630859e-03, -5.12695312e-03,
            9.49859619e-04, -7.01904297e-03, -3.87573242e-03, 2.99072266e-03,
            -3.84521484e-03, 6.10351562e-03, 9.33837891e-03, -3.02124023e-03,
            9.76562500e-03, -1.70898438e-03, -1.19018555e-02, 1.32446289e-02,
            1.00708008e-02, -6.53076172e-03, 1.55639648e-02, 1.57470703e-02,
            8.23974609e-03, -2.92968750e-03, -1.98974609e-02, 3.58581543e-03,
            -5.21850586e-03, -4.91333008e-03, 3.29589844e-03, 5.15747070e-03,
            9.70458984e-03, -3.40270996e-03, -1.46484375e-02, 2.54821777e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_2 = np.array(
        [[[[4.94384766e-03, -1.40380859e-03, 6.53076172e-03, 1.51977539e-02,
            -2.33459473e-03, -2.34985352e-03, 1.56250000e-02, 7.69042969e-03,
            -4.69970703e-03, 5.43212891e-03, -4.73022461e-03, -4.66918945e-03,
            2.44140625e-03, -1.91650391e-02, -1.72119141e-02, -5.67626953e-03,
            -1.01318359e-02, 3.26538086e-03, -8.91113281e-03, -1.41601562e-02,
            1.46484375e-02, -2.39562988e-03, 8.73565674e-04, -1.42211914e-02,
            -5.43212891e-03, 9.84191895e-04, -1.14746094e-02, 3.81469727e-03,
            -6.01196289e-03, -2.88391113e-03, -5.73730469e-03, 1.85546875e-02,
            -2.47955322e-04, -1.05590820e-02, 8.11767578e-03, -1.20849609e-02,
            1.93786621e-03, -1.96533203e-02, -1.33056641e-02, 2.13623047e-03,
            7.38525391e-03, 1.81579590e-03, -1.20544434e-03, -3.00598145e-03,
            -1.47705078e-02, -7.32421875e-03, -4.48608398e-03, 1.06201172e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.25012207e-03, -3.82995605e-03, -6.77490234e-03, 6.07299805e-03,
            1.03759766e-02, 9.33837891e-03, -8.36181641e-03, -3.03649902e-03,
            3.12805176e-03, 9.76562500e-03, -4.82177734e-03, -1.88446045e-03,
            -1.11694336e-02, -1.20239258e-02, 7.93457031e-03, 1.34887695e-02,
            -8.16345215e-04, 1.00097656e-02, 3.66210938e-03, -6.37817383e-03,
            3.64685059e-03, 1.53808594e-02, -2.19345093e-04, 1.56250000e-02,
            -2.62451172e-02, 8.17871094e-03, 8.77380371e-04, -2.99072266e-03,
            9.23156738e-04, -1.98974609e-02, -2.18200684e-03, 3.50952148e-03,
            1.48315430e-02, -5.31005859e-03, -8.11767578e-03, -5.03540039e-03,
            9.15527344e-03, 3.34167480e-03, -5.34057617e-03, 5.12695312e-03,
            1.05285645e-03, 9.70458984e-03, -7.01904297e-03, -3.34167480e-03,
            -3.87573242e-03, -1.45874023e-02, 2.96020508e-03, 2.62451172e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_3 = np.array(
        [[[[4.94384766e-03, 6.53076172e-03, -2.34985352e-03, 1.58691406e-02,
            -4.69970703e-03, -4.60815430e-03, 2.79235840e-03, -1.72119141e-02,
            -1.01318359e-02, -9.09423828e-03, 1.45874023e-02, 6.71386719e-04,
            -5.43212891e-03, -1.16577148e-02, -6.01196289e-03, -6.07299805e-03,
            -1.39236450e-04, 8.11767578e-03, 2.31933594e-03, -1.33056641e-02,
            7.38525391e-03, -1.19018555e-03, -1.46484375e-02, -4.45556641e-03,
            -1.35040283e-03, 1.53808594e-02, -2.30407715e-03, 7.56835938e-03,
            5.43212891e-03, -4.63867188e-03, -1.91650391e-02, -5.61523438e-03,
            3.25012207e-03, -1.41601562e-02, -2.39562988e-03, -1.42211914e-02,
            1.15203857e-03, 3.78417969e-03, -2.96020508e-03, 1.86767578e-02,
            -1.05590820e-02, -1.20849609e-02, -1.96533203e-02, 1.94549561e-03,
            1.73950195e-03, -3.02124023e-03, -7.01904297e-03, 1.06201172e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.25012207e-03, -6.80541992e-03, 1.03149414e-02, -8.36181641e-03,
            3.37219238e-03, -4.79125977e-03, -1.09863281e-02, 8.23974609e-03,
            -9.07897949e-04, 3.64685059e-03, 3.79943848e-03, -4.61578369e-04,
            -2.61230469e-02, 8.92639160e-04, 1.15203857e-03, -2.16674805e-03,
            1.48315430e-02, -8.05664062e-03, 9.15527344e-03, -5.40161133e-03,
            9.91821289e-04, -7.04956055e-03, -4.02832031e-03, 2.97546387e-03,
            -3.84521484e-03, 6.07299805e-03, 9.33837891e-03, -3.17382812e-03,
            9.76562500e-03, -1.92260742e-03, -1.19628906e-02, 1.37939453e-02,
            1.00097656e-02, -6.46972656e-03, 1.53198242e-02, 1.56250000e-02,
            8.30078125e-03, -2.99072266e-03, -1.98974609e-02, 3.57055664e-03,
            -5.31005859e-03, -4.88281250e-03, 3.23486328e-03, 5.09643555e-03,
            9.70458984e-03, -3.18908691e-03, -1.46484375e-02, 2.62451172e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_4 = np.array(
        [[[[4.88281250e-03, -1.37329102e-03, 6.43920898e-03, 1.52587891e-02,
            -2.34985352e-03, -2.31933594e-03, 1.58691406e-02, 7.62939453e-03,
            -4.69970703e-03, 5.43212891e-03, -4.48608398e-03, -4.66918945e-03,
            2.42614746e-03, -1.86767578e-02, -1.72119141e-02, -5.61523438e-03,
            -1.01318359e-02, 3.15856934e-03, -9.09423828e-03, -1.40991211e-02,
            1.47705078e-02, -2.31933594e-03, 7.40051270e-04, -1.40991211e-02,
            -5.40161133e-03, 1.07574463e-03, -1.15966797e-02, 3.66210938e-03,
            -6.01196289e-03, -2.89916992e-03, -6.25610352e-03, 1.85546875e-02,
            -8.58306885e-05, -1.05590820e-02, 8.30078125e-03, -1.22680664e-02,
            2.07519531e-03, -1.95312500e-02, -1.34277344e-02, 2.04467773e-03,
            7.35473633e-03, 1.75476074e-03, -1.01470947e-03, -3.03649902e-03,
            -1.47094727e-02, -7.23266602e-03, -4.60815430e-03, 1.07421875e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.37219238e-03, -3.87573242e-03, -6.77490234e-03, 6.10351562e-03,
            1.03149414e-02, 9.39941406e-03, -8.36181641e-03, -3.06701660e-03,
            3.03649902e-03, 9.76562500e-03, -4.88281250e-03, -1.82342529e-03,
            -1.09863281e-02, -1.19628906e-02, 8.17871094e-03, 1.36108398e-02,
            -7.28607178e-04, 9.94873047e-03, 3.67736816e-03, -6.37817383e-03,
            3.60107422e-03, 1.52587891e-02, -3.31878662e-04, 1.56250000e-02,
            -2.62451172e-02, 8.23974609e-03, 8.50677490e-04, -2.94494629e-03,
            9.19342041e-04, -1.97753906e-02, -2.18200684e-03, 3.47900391e-03,
            1.47705078e-02, -5.27954102e-03, -7.99560547e-03, -5.00488281e-03,
            9.15527344e-03, 3.18908691e-03, -5.27954102e-03, 5.12695312e-03,
            9.76562500e-04, 9.58251953e-03, -7.04956055e-03, -3.34167480e-03,
            -3.89099121e-03, -1.48315430e-02, 2.96020508e-03, 2.67028809e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    return {
        "output_1": output_1,
        "output_2": output_2,
        "output_3": output_3,
        "output_4": output_4
    }

def get_gpu_datas() -> dict[str, np.ndarray]:
    """Generate gpu data for test."""
    output_1 = np.array(
        [[[[5.06591797e-03, 6.43920898e-03, -2.27355957e-03, 1.57470703e-02,
            -4.69970703e-03, -4.66918945e-03, 2.74658203e-03, -1.72119141e-02,
            -1.01318359e-02, -9.09423828e-03, 1.46484375e-02, 6.71386719e-04,
            -5.43212891e-03, -1.16577148e-02, -6.01196289e-03, -6.04248047e-03,
            -1.32560730e-04, 8.23974609e-03, 2.24304199e-03, -1.33056641e-02,
            7.38525391e-03, -1.22833252e-03, -1.47094727e-02, -4.39453125e-03,
            -1.47247314e-03, 1.52587891e-02, -2.25830078e-03, 7.65991211e-03,
            5.43212891e-03, -4.69970703e-03, -1.92871094e-02, -5.61523438e-03,
            2.99072266e-03, -1.40991211e-02, -2.39562988e-03, -1.42822266e-02,
            1.13677979e-03, 3.75366211e-03, -3.06701660e-03, 1.85546875e-02,
            -1.05590820e-02, -1.23291016e-02, -1.96533203e-02, 1.93023682e-03,
            1.72424316e-03, -2.99072266e-03, -7.08007812e-03, 1.06811523e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.26538086e-03, -6.77490234e-03, 1.03149414e-02, -8.30078125e-03,
            3.28063965e-03, -4.79125977e-03, -1.10473633e-02, 8.17871094e-03,
            -7.89642334e-04, 3.61633301e-03, 3.61633301e-03, -3.35693359e-04,
            -2.62451172e-02, 6.63757324e-04, 9.46044922e-04, -2.19726562e-03,
            1.48315430e-02, -7.99560547e-03, 9.21630859e-03, -5.12695312e-03,
            9.49859619e-04, -7.01904297e-03, -3.87573242e-03, 2.99072266e-03,
            -3.84521484e-03, 6.10351562e-03, 9.33837891e-03, -3.02124023e-03,
            9.76562500e-03, -1.70898438e-03, -1.19018555e-02, 1.32446289e-02,
            1.00708008e-02, -6.53076172e-03, 1.55639648e-02, 1.57470703e-02,
            8.23974609e-03, -2.92968750e-03, -1.98974609e-02, 3.58581543e-03,
            -5.21850586e-03, -4.91333008e-03, 3.29589844e-03, 5.15747070e-03,
            9.70458984e-03, -3.40270996e-03, -1.46484375e-02, 2.54821777e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_2 = np.array(
        [[[[4.94384766e-03, -1.40380859e-03, 6.53076172e-03, 1.51977539e-02,
            -2.33459473e-03, -2.34985352e-03, 1.56250000e-02, 7.69042969e-03,
            -4.69970703e-03, 5.43212891e-03, -4.73022461e-03, -4.66918945e-03,
            2.44140625e-03, -1.91650391e-02, -1.72119141e-02, -5.67626953e-03,
            -1.01318359e-02, 3.26538086e-03, -8.91113281e-03, -1.41601562e-02,
            1.46484375e-02, -2.39562988e-03, 8.73565674e-04, -1.42211914e-02,
            -5.43212891e-03, 9.84191895e-04, -1.14746094e-02, 3.81469727e-03,
            -6.01196289e-03, -2.88391113e-03, -5.73730469e-03, 1.85546875e-02,
            -2.47955322e-04, -1.05590820e-02, 8.11767578e-03, -1.20849609e-02,
            1.93786621e-03, -1.96533203e-02, -1.33056641e-02, 2.13623047e-03,
            7.38525391e-03, 1.81579590e-03, -1.20544434e-03, -3.00598145e-03,
            -1.47705078e-02, -7.32421875e-03, -4.48608398e-03, 1.06201172e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.25012207e-03, -3.82995605e-03, -6.77490234e-03, 6.07299805e-03,
            1.03759766e-02, 9.33837891e-03, -8.36181641e-03, -3.03649902e-03,
            3.12805176e-03, 9.76562500e-03, -4.82177734e-03, -1.88446045e-03,
            -1.11694336e-02, -1.20239258e-02, 7.93457031e-03, 1.34887695e-02,
            -8.16345215e-04, 1.00097656e-02, 3.66210938e-03, -6.37817383e-03,
            3.64685059e-03, 1.53808594e-02, -2.19345093e-04, 1.56250000e-02,
            -2.62451172e-02, 8.17871094e-03, 8.77380371e-04, -2.99072266e-03,
            9.23156738e-04, -1.98974609e-02, -2.18200684e-03, 3.50952148e-03,
            1.48315430e-02, -5.31005859e-03, -8.11767578e-03, -5.03540039e-03,
            9.15527344e-03, 3.34167480e-03, -5.34057617e-03, 5.12695312e-03,
            1.05285645e-03, 9.70458984e-03, -7.01904297e-03, -3.34167480e-03,
            -3.87573242e-03, -1.45874023e-02, 2.96020508e-03, 2.62451172e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_3 = np.array(
        [[[[4.94384766e-03, 6.53076172e-03, -2.34985352e-03, 1.58691406e-02,
            -4.69970703e-03, -4.60815430e-03, 2.79235840e-03, -1.72119141e-02,
            -1.01318359e-02, -9.09423828e-03, 1.45874023e-02, 6.71386719e-04,
            -5.43212891e-03, -1.16577148e-02, -6.01196289e-03, -6.07299805e-03,
            -1.39236450e-04, 8.11767578e-03, 2.31933594e-03, -1.33056641e-02,
            7.38525391e-03, -1.19018555e-03, -1.46484375e-02, -4.45556641e-03,
            -1.35040283e-03, 1.53808594e-02, -2.30407715e-03, 7.56835938e-03,
            5.43212891e-03, -4.63867188e-03, -1.91650391e-02, -5.61523438e-03,
            3.25012207e-03, -1.41601562e-02, -2.39562988e-03, -1.42211914e-02,
            1.15203857e-03, 3.78417969e-03, -2.96020508e-03, 1.86767578e-02,
            -1.05590820e-02, -1.20849609e-02, -1.96533203e-02, 1.94549561e-03,
            1.73950195e-03, -3.02124023e-03, -7.01904297e-03, 1.06201172e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.25012207e-03, -6.80541992e-03, 1.03149414e-02, -8.36181641e-03,
            3.37219238e-03, -4.79125977e-03, -1.09863281e-02, 8.23974609e-03,
            -9.07897949e-04, 3.64685059e-03, 3.79943848e-03, -4.61578369e-04,
            -2.61230469e-02, 8.92639160e-04, 1.15203857e-03, -2.16674805e-03,
            1.48315430e-02, -8.05664062e-03, 9.15527344e-03, -5.40161133e-03,
            9.91821289e-04, -7.04956055e-03, -4.02832031e-03, 2.97546387e-03,
            -3.84521484e-03, 6.07299805e-03, 9.33837891e-03, -3.17382812e-03,
            9.76562500e-03, -1.92260742e-03, -1.19628906e-02, 1.37939453e-02,
            1.00097656e-02, -6.46972656e-03, 1.53198242e-02, 1.56250000e-02,
            8.30078125e-03, -2.99072266e-03, -1.98974609e-02, 3.57055664e-03,
            -5.31005859e-03, -4.88281250e-03, 3.23486328e-03, 5.09643555e-03,
            9.70458984e-03, -3.18908691e-03, -1.46484375e-02, 2.62451172e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    output_4 = np.array(
        [[[[4.88281250e-03, -1.37329102e-03, 6.43920898e-03, 1.52587891e-02,
            -2.34985352e-03, -2.31933594e-03, 1.58691406e-02, 7.62939453e-03,
            -4.69970703e-03, 5.43212891e-03, -4.48608398e-03, -4.66918945e-03,
            2.42614746e-03, -1.86767578e-02, -1.72119141e-02, -5.61523438e-03,
            -1.01318359e-02, 3.15856934e-03, -9.09423828e-03, -1.40991211e-02,
            1.47705078e-02, -2.31933594e-03, 7.40051270e-04, -1.40991211e-02,
            -5.40161133e-03, 1.07574463e-03, -1.15966797e-02, 3.66210938e-03,
            -6.01196289e-03, -2.89916992e-03, -6.25610352e-03, 1.85546875e-02,
            -8.58306885e-05, -1.05590820e-02, 8.30078125e-03, -1.22680664e-02,
            2.07519531e-03, -1.95312500e-02, -1.34277344e-02, 2.04467773e-03,
            7.35473633e-03, 1.75476074e-03, -1.01470947e-03, -3.03649902e-03,
            -1.47094727e-02, -7.23266602e-03, -4.60815430e-03, 1.07421875e-02,
            3.43322754e-03, -1.75781250e-02],
           [3.37219238e-03, -3.87573242e-03, -6.77490234e-03, 6.10351562e-03,
            1.03149414e-02, 9.39941406e-03, -8.36181641e-03, -3.06701660e-03,
            3.03649902e-03, 9.76562500e-03, -4.88281250e-03, -1.82342529e-03,
            -1.09863281e-02, -1.19628906e-02, 8.17871094e-03, 1.36108398e-02,
            -7.28607178e-04, 9.94873047e-03, 3.67736816e-03, -6.37817383e-03,
            3.60107422e-03, 1.52587891e-02, -3.31878662e-04, 1.56250000e-02,
            -2.62451172e-02, 8.23974609e-03, 8.50677490e-04, -2.94494629e-03,
            9.19342041e-04, -1.97753906e-02, -2.18200684e-03, 3.47900391e-03,
            1.47705078e-02, -5.27954102e-03, -7.99560547e-03, -5.00488281e-03,
            9.15527344e-03, 3.18908691e-03, -5.27954102e-03, 5.12695312e-03,
            9.76562500e-04, 9.58251953e-03, -7.04956055e-03, -3.34167480e-03,
            -3.89099121e-03, -1.48315430e-02, 2.96020508e-03, 2.67028809e-03,
            5.10215759e-05, -2.34985352e-03]]]])
    return {
        "output_1": output_1,
        "output_2": output_2,
        "output_3": output_3,
        "output_4": output_4,
    }

GOLDEN_DATA = get_golden()
GPU_DATA = get_gpu_datas()
